import idc
import idaapi

class RegisterXrefs(object):

	UP   = 'Up  '
	DOWN = 'Down'
	THIS = '-   '

	READ    = 'r'
	WRITE   = 'w'
	EXECUTE = 'x'

	JUMP_BRANCH_INSTRUCTIONS = ['j', 'b']

	def __init__(self):
		self.xrefs = {}
		self.function = ''
		self._profile_function()

	def _profile_function(self):
		current_ea = ScreenEA()
		current_function = idc.GetFunctionName(current_ea)
		current_function_ea = idc.LocByName(current_function)

		if current_function:
			self.function = current_function

		ea = start_ea = idc.GetFunctionAttr(current_function_ea,  idc.FUNCATTR_START)
		end_ea = idc.GetFunctionAttr(current_function_ea, idc.FUNCATTR_END)

		self.highlighted = idaapi.get_highlighted_identifier()

		while ea < end_ea and ea != idc.BADADDR and self.highlighted:

			i = 0
			match = False
			comment = None
			opnds = []
			mnem = idc.GetMnem(ea)

			if self.highlighted in mnem:
				match = True

			if idaapi.is_call_insn(ea):
				for xref in idautils.XrefsFrom(ea):
					if xref.type != 21:
						name = idc.Name(xref.to)
						if name and self.highlighted in name:
							match = True
							break
	
			while True:
				opnd = idc.GetOpnd(ea, i)
				if opnd:
					if self.highlighted in opnd:
						match = True

					opnds.append(opnd)
					i += 1
				else:
					break

			if not match:
				comment = idc.GetCommentEx(ea, 0)
				if comment and self.highlighted in comment:
					match = True
				else:
					comment = idc.GetCommentEx(ea, 1)
					if comment and self.highlighted in comment:
						match = True
					else:
						comment = None

			if match:
				if ea > current_ea:
					direction = self.DOWN
				elif ea < current_ea:
					direction = self.UP
				else:
					direction = self.THIS

				optype = self.operation_type(self.highlighted, mnem, opnds, comment)
				if not optype:
					optype = self.READ

				self.xrefs[ea] = {
					'offset' 	: idc.GetFuncOffset(ea),
					'mnem'	 	: mnem,
					'operands'	: opnds,
					'type'		: optype,
					'direction'	: direction,
					'text'		: idc.GetDisasm(ea),
				}

			ea += idaapi.decode_insn(ea)
	
	def operation_type(self, highlighted, instruction, operands, comment):

		if instruction[0:1].lower() in self.JUMP_BRANCH_INSTRUCTIONS:

			return self.EXECUTE

		elif instruction == highlighted:
		
			return self.EXECUTE
		
		# This is generally, but not universally true
		elif operands[0] == highlighted:

			return self.WRITE
		
		return self.READ


class localizedxrefs_t(idaapi.plugin_t):
	flags = 0
	comment = "IDA Localized Xrefs"
	help = ""
	wanted_name = "Localized Xrefs"
	wanted_hotkey = ""

	DELIM = '-' * 86
	HEADER = '\nXrefs to %s from %s:'

	# Add RegsiterXrefs sub-classes here, if needed
	supported_architectures = {
		'default'	: RegisterXrefs,
	}

	def init(self):
		self.menu_context = idaapi.add_menu_item("Jump/", "Find localized xrefs", "Alt-8", 0, self.run, (None,))
		return idaapi.PLUGIN_KEEP

	def term(self):
		idaapi.del_menu_item(self.menu_context)
		return None

	def run(self, arg):
		for (cpu, subclass) in self.supported_architectures.iteritems():
			if idc.GetShortPrm(idc.INF_PROCNAME).lower().startswith(cpu.lower()):
				break
			else:
				subclass = None

		if not subclass:
			subclass = supported_architectures['default']

		if callable(subclass):
			fmt = ''

			r = subclass()
			offsets = r.xrefs.keys()
			offsets.sort()

			if r.highlighted:
				print self.HEADER % (r.highlighted, r.function)
				print self.DELIM
			
				for ea in offsets:
					info = r.xrefs[ea]
	
					if not fmt:
						fmt = "%%s   %%s   %%-%ds   %%s" % (len(info['offset']) + 15)

					print fmt % (info['direction'], info['type'], info['offset'], info['text'])
	
				print self.DELIM

def PLUGIN_ENTRY():
	return localizedxrefs_t()

